{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train a character-level GPT on some text data\n",
    "\n",
    "The inputs here are simple text files, which we chop up to individual characters and then train GPT on. So you could say this is a char-transformer instead of a char-rnn. Doesn't quite roll off the tongue as well. In this example we will feed it some shakespear, which we'll get it to predict character-level."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up logging\n",
    "import logging\n",
    "logging.basicConfig(\n",
    "        format=\"%(asctime)s - %(levelname)s - %(name)s -   %(message)s\",\n",
    "        datefmt=\"%m/%d/%Y %H:%M:%S\",\n",
    "        level=logging.INFO,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make deterministic\n",
    "from mingpt.utils import set_seed\n",
    "set_seed(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "class CharDataset(Dataset):\n",
    "\n",
    "    def __init__(self, data, block_size):\n",
    "        chars = list(set(data))\n",
    "        data_size, vocab_size = len(data), len(chars)\n",
    "        print('data has %d characters, %d unique.' % (data_size, vocab_size))\n",
    "        \n",
    "        self.stoi = { ch:i for i,ch in enumerate(chars) }\n",
    "        self.itos = { i:ch for i,ch in enumerate(chars) }\n",
    "        self.block_size = block_size\n",
    "        self.vocab_size = vocab_size\n",
    "        self.data = data\n",
    "    \n",
    "    def __len__(self):\n",
    "        return math.ceil(len(self.data) / (self.block_size + 1))\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        # we're actually going to \"cheat\" and pick a spot in the dataset at random\n",
    "        i = np.random.randint(0, len(self.data) - (self.block_size + 1))\n",
    "        chunk = self.data[i:i+self.block_size+1]\n",
    "        dix = [self.stoi[s] for s in chunk]\n",
    "        x = torch.tensor(dix[:-1], dtype=torch.long)\n",
    "        y = torch.tensor(dix[1:], dtype=torch.long)\n",
    "        return x, y\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "block_size = 128 # spatial extent of the model for its context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "data has 1115394 characters, 65 unique.\n"
     ]
    }
   ],
   "source": [
    "# you can download this file at https://github.com/karpathy/char-rnn/blob/master/data/tinyshakespeare/input.txt\n",
    "text = open('input.txt', 'r').read() # don't worry we won't run out of file handles\n",
    "train_dataset = CharDataset(text, block_size) # one line of poem is roughly 50 characters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "08/17/2020 00:11:58 - INFO - mingpt.model -   number of parameters: 2.535219e+07\n"
     ]
    }
   ],
   "source": [
    "from mingpt.model import GPT, GPTConfig\n",
    "mconf = GPTConfig(train_dataset.vocab_size, train_dataset.block_size,\n",
    "                  n_layer=8, n_head=8, n_embd=512)\n",
    "model = GPT(mconf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/17 [00:00<?, ?it/s]/apcv/shared/conda-envs/apcv-6244e1d-566/lib/python3.8/site-packages/torch/nn/parallel/_functions.py:61: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      "epoch 1 iter 16: train loss 3.31022. lr 5.999637e-04: 100%|██████████| 17/17 [00:36<00:00,  2.18s/it]\n",
      "epoch 2 iter 16: train loss 2.89320. lr 5.998533e-04: 100%|██████████| 17/17 [00:04<00:00,  3.78it/s]\n",
      "epoch 3 iter 16: train loss 2.63845. lr 5.996690e-04: 100%|██████████| 17/17 [00:04<00:00,  3.74it/s]\n",
      "epoch 4 iter 16: train loss 2.54588. lr 5.994107e-04: 100%|██████████| 17/17 [00:04<00:00,  3.87it/s]\n",
      "epoch 5 iter 16: train loss 2.49512. lr 5.990785e-04: 100%|██████████| 17/17 [00:04<00:00,  3.98it/s]\n",
      "epoch 6 iter 16: train loss 2.46732. lr 5.986726e-04: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 7 iter 16: train loss 2.44716. lr 5.981929e-04: 100%|██████████| 17/17 [00:04<00:00,  3.95it/s]\n",
      "epoch 8 iter 16: train loss 2.37363. lr 5.976397e-04: 100%|██████████| 17/17 [00:04<00:00,  3.93it/s]\n",
      "epoch 9 iter 16: train loss 2.34669. lr 5.970130e-04: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 10 iter 16: train loss 2.28792. lr 5.963130e-04: 100%|██████████| 17/17 [00:04<00:00,  3.98it/s]\n",
      "epoch 11 iter 16: train loss 2.21925. lr 5.955399e-04: 100%|██████████| 17/17 [00:04<00:00,  4.04it/s]\n",
      "epoch 12 iter 16: train loss 2.16131. lr 5.946939e-04: 100%|██████████| 17/17 [00:04<00:00,  4.04it/s]\n",
      "epoch 13 iter 16: train loss 2.12197. lr 5.937751e-04: 100%|██████████| 17/17 [00:04<00:00,  4.10it/s]\n",
      "epoch 14 iter 16: train loss 2.06564. lr 5.927839e-04: 100%|██████████| 17/17 [00:04<00:00,  4.01it/s]\n",
      "epoch 15 iter 16: train loss 2.00401. lr 5.917204e-04: 100%|██████████| 17/17 [00:04<00:00,  3.97it/s]\n",
      "epoch 16 iter 16: train loss 1.96109. lr 5.905849e-04: 100%|██████████| 17/17 [00:04<00:00,  3.97it/s]\n",
      "epoch 17 iter 16: train loss 1.89554. lr 5.893777e-04: 100%|██████████| 17/17 [00:04<00:00,  3.97it/s]\n",
      "epoch 18 iter 16: train loss 1.85840. lr 5.880992e-04: 100%|██████████| 17/17 [00:04<00:00,  4.04it/s]\n",
      "epoch 19 iter 16: train loss 1.80772. lr 5.867495e-04: 100%|██████████| 17/17 [00:04<00:00,  3.99it/s]\n",
      "epoch 20 iter 16: train loss 1.76782. lr 5.853291e-04: 100%|██████████| 17/17 [00:04<00:00,  4.06it/s]\n",
      "epoch 21 iter 16: train loss 1.73638. lr 5.838382e-04: 100%|██████████| 17/17 [00:04<00:00,  4.05it/s]\n",
      "epoch 22 iter 16: train loss 1.71822. lr 5.822774e-04: 100%|██████████| 17/17 [00:04<00:00,  3.99it/s]\n",
      "epoch 23 iter 16: train loss 1.65840. lr 5.806468e-04: 100%|██████████| 17/17 [00:04<00:00,  4.07it/s]\n",
      "epoch 24 iter 16: train loss 1.62377. lr 5.789471e-04: 100%|██████████| 17/17 [00:04<00:00,  4.07it/s]\n",
      "epoch 25 iter 16: train loss 1.58972. lr 5.771785e-04: 100%|██████████| 17/17 [00:04<00:00,  4.11it/s]\n",
      "epoch 26 iter 16: train loss 1.56834. lr 5.753415e-04: 100%|██████████| 17/17 [00:04<00:00,  4.08it/s]\n",
      "epoch 27 iter 16: train loss 1.53979. lr 5.734365e-04: 100%|██████████| 17/17 [00:04<00:00,  4.03it/s]\n",
      "epoch 28 iter 16: train loss 1.50191. lr 5.714641e-04: 100%|██████████| 17/17 [00:04<00:00,  3.98it/s]\n",
      "epoch 29 iter 16: train loss 1.48912. lr 5.694247e-04: 100%|██████████| 17/17 [00:04<00:00,  3.99it/s]\n",
      "epoch 30 iter 16: train loss 1.46471. lr 5.673188e-04: 100%|██████████| 17/17 [00:04<00:00,  3.99it/s]\n",
      "epoch 31 iter 16: train loss 1.42849. lr 5.651469e-04: 100%|██████████| 17/17 [00:04<00:00,  3.99it/s]\n",
      "epoch 32 iter 16: train loss 1.43040. lr 5.629096e-04: 100%|██████████| 17/17 [00:04<00:00,  3.99it/s]\n",
      "epoch 33 iter 16: train loss 1.41800. lr 5.606075e-04: 100%|██████████| 17/17 [00:04<00:00,  3.98it/s]\n",
      "epoch 34 iter 16: train loss 1.38389. lr 5.582410e-04: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 35 iter 16: train loss 1.39671. lr 5.558108e-04: 100%|██████████| 17/17 [00:04<00:00,  3.97it/s]\n",
      "epoch 36 iter 16: train loss 1.37993. lr 5.533175e-04: 100%|██████████| 17/17 [00:04<00:00,  3.94it/s]\n",
      "epoch 37 iter 16: train loss 1.34003. lr 5.507617e-04: 100%|██████████| 17/17 [00:04<00:00,  3.98it/s]\n",
      "epoch 38 iter 16: train loss 1.33921. lr 5.481440e-04: 100%|██████████| 17/17 [00:04<00:00,  3.93it/s]\n",
      "epoch 39 iter 16: train loss 1.33100. lr 5.454651e-04: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 40 iter 16: train loss 1.31935. lr 5.427256e-04: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 41 iter 16: train loss 1.30351. lr 5.399262e-04: 100%|██████████| 17/17 [00:04<00:00,  3.97it/s]\n",
      "epoch 42 iter 16: train loss 1.29172. lr 5.370676e-04: 100%|██████████| 17/17 [00:04<00:00,  3.95it/s]\n",
      "epoch 43 iter 16: train loss 1.28822. lr 5.341505e-04: 100%|██████████| 17/17 [00:04<00:00,  3.98it/s]\n",
      "epoch 44 iter 16: train loss 1.26141. lr 5.311756e-04: 100%|██████████| 17/17 [00:04<00:00,  3.98it/s]\n",
      "epoch 45 iter 16: train loss 1.26616. lr 5.281437e-04: 100%|██████████| 17/17 [00:04<00:00,  4.00it/s]\n",
      "epoch 46 iter 16: train loss 1.25256. lr 5.250555e-04: 100%|██████████| 17/17 [00:04<00:00,  3.98it/s]\n",
      "epoch 47 iter 16: train loss 1.22032. lr 5.219118e-04: 100%|██████████| 17/17 [00:04<00:00,  3.97it/s]\n",
      "epoch 48 iter 16: train loss 1.23235. lr 5.187133e-04: 100%|██████████| 17/17 [00:04<00:00,  3.94it/s]\n",
      "epoch 49 iter 16: train loss 1.23283. lr 5.154608e-04: 100%|██████████| 17/17 [00:04<00:00,  3.95it/s]\n",
      "epoch 50 iter 16: train loss 1.20031. lr 5.121552e-04: 100%|██████████| 17/17 [00:04<00:00,  3.98it/s]\n",
      "epoch 51 iter 16: train loss 1.18663. lr 5.087972e-04: 100%|██████████| 17/17 [00:04<00:00,  3.99it/s]\n",
      "epoch 52 iter 16: train loss 1.19119. lr 5.053876e-04: 100%|██████████| 17/17 [00:04<00:00,  4.02it/s]\n",
      "epoch 53 iter 16: train loss 1.19220. lr 5.019275e-04: 100%|██████████| 17/17 [00:04<00:00,  4.03it/s]\n",
      "epoch 54 iter 16: train loss 1.16783. lr 4.984174e-04: 100%|██████████| 17/17 [00:04<00:00,  4.00it/s]\n",
      "epoch 55 iter 16: train loss 1.15194. lr 4.948584e-04: 100%|██████████| 17/17 [00:04<00:00,  4.03it/s]\n",
      "epoch 56 iter 16: train loss 1.14275. lr 4.912514e-04: 100%|██████████| 17/17 [00:04<00:00,  4.04it/s]\n",
      "epoch 57 iter 16: train loss 1.14208. lr 4.875971e-04: 100%|██████████| 17/17 [00:04<00:00,  3.99it/s]\n",
      "epoch 58 iter 16: train loss 1.13699. lr 4.838966e-04: 100%|██████████| 17/17 [00:04<00:00,  3.98it/s]\n",
      "epoch 59 iter 16: train loss 1.12777. lr 4.801506e-04: 100%|██████████| 17/17 [00:04<00:00,  3.99it/s]\n",
      "epoch 60 iter 16: train loss 1.11677. lr 4.763603e-04: 100%|██████████| 17/17 [00:04<00:00,  3.94it/s]\n",
      "epoch 61 iter 16: train loss 1.11154. lr 4.725264e-04: 100%|██████████| 17/17 [00:04<00:00,  4.01it/s]\n",
      "epoch 62 iter 16: train loss 1.09913. lr 4.686499e-04: 100%|██████████| 17/17 [00:04<00:00,  4.03it/s]\n",
      "epoch 63 iter 16: train loss 1.08472. lr 4.647318e-04: 100%|██████████| 17/17 [00:04<00:00,  4.12it/s]\n",
      "epoch 64 iter 16: train loss 1.07229. lr 4.607731e-04: 100%|██████████| 17/17 [00:04<00:00,  4.08it/s]\n",
      "epoch 65 iter 16: train loss 1.05527. lr 4.567747e-04: 100%|██████████| 17/17 [00:04<00:00,  4.03it/s]\n",
      "epoch 66 iter 16: train loss 1.05399. lr 4.527376e-04: 100%|██████████| 17/17 [00:04<00:00,  4.09it/s]\n",
      "epoch 67 iter 16: train loss 1.04091. lr 4.486628e-04: 100%|██████████| 17/17 [00:04<00:00,  4.12it/s]\n",
      "epoch 68 iter 16: train loss 1.02974. lr 4.445513e-04: 100%|██████████| 17/17 [00:04<00:00,  4.16it/s]\n",
      "epoch 69 iter 16: train loss 1.02136. lr 4.404042e-04: 100%|██████████| 17/17 [00:04<00:00,  4.12it/s]\n",
      "epoch 70 iter 16: train loss 0.99327. lr 4.362224e-04: 100%|██████████| 17/17 [00:04<00:00,  4.04it/s]\n",
      "epoch 71 iter 16: train loss 1.00298. lr 4.320070e-04: 100%|██████████| 17/17 [00:04<00:00,  4.08it/s]\n",
      "epoch 72 iter 16: train loss 0.98914. lr 4.277590e-04: 100%|██████████| 17/17 [00:04<00:00,  4.05it/s]\n",
      "epoch 73 iter 16: train loss 0.97867. lr 4.234795e-04: 100%|██████████| 17/17 [00:04<00:00,  4.07it/s]\n",
      "epoch 74 iter 16: train loss 0.96466. lr 4.191696e-04: 100%|██████████| 17/17 [00:04<00:00,  4.05it/s]\n",
      "epoch 75 iter 16: train loss 0.95394. lr 4.148302e-04: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 76 iter 16: train loss 0.94508. lr 4.104625e-04: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 77 iter 16: train loss 0.92286. lr 4.060675e-04: 100%|██████████| 17/17 [00:04<00:00,  3.95it/s]\n",
      "epoch 78 iter 16: train loss 0.93483. lr 4.016464e-04: 100%|██████████| 17/17 [00:04<00:00,  3.92it/s]\n",
      "epoch 79 iter 16: train loss 0.91154. lr 3.972002e-04: 100%|██████████| 17/17 [00:04<00:00,  3.93it/s]\n",
      "epoch 80 iter 16: train loss 0.90103. lr 3.927300e-04: 100%|██████████| 17/17 [00:04<00:00,  3.99it/s]\n",
      "epoch 81 iter 16: train loss 0.88080. lr 3.882369e-04: 100%|██████████| 17/17 [00:04<00:00,  3.98it/s]\n",
      "epoch 82 iter 16: train loss 0.87298. lr 3.837220e-04: 100%|██████████| 17/17 [00:04<00:00,  3.97it/s]\n",
      "epoch 83 iter 16: train loss 0.87236. lr 3.791865e-04: 100%|██████████| 17/17 [00:04<00:00,  3.94it/s]\n",
      "epoch 84 iter 16: train loss 0.84863. lr 3.746315e-04: 100%|██████████| 17/17 [00:04<00:00,  3.98it/s]\n",
      "epoch 85 iter 16: train loss 0.84596. lr 3.700580e-04: 100%|██████████| 17/17 [00:04<00:00,  3.92it/s]\n",
      "epoch 86 iter 16: train loss 0.83594. lr 3.654672e-04: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 87 iter 16: train loss 0.80802. lr 3.608603e-04: 100%|██████████| 17/17 [00:04<00:00,  3.99it/s]\n",
      "epoch 88 iter 16: train loss 0.81852. lr 3.562384e-04: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 89 iter 16: train loss 0.79750. lr 3.516026e-04: 100%|██████████| 17/17 [00:04<00:00,  3.95it/s]\n",
      "epoch 90 iter 16: train loss 0.79546. lr 3.469540e-04: 100%|██████████| 17/17 [00:04<00:00,  3.94it/s]\n",
      "epoch 91 iter 16: train loss 0.78435. lr 3.422939e-04: 100%|██████████| 17/17 [00:04<00:00,  3.94it/s]\n",
      "epoch 92 iter 16: train loss 0.77504. lr 3.376233e-04: 100%|██████████| 17/17 [00:04<00:00,  3.93it/s]\n",
      "epoch 93 iter 16: train loss 0.75585. lr 3.329435e-04: 100%|██████████| 17/17 [00:04<00:00,  4.03it/s]\n",
      "epoch 94 iter 16: train loss 0.74559. lr 3.282555e-04: 100%|██████████| 17/17 [00:04<00:00,  4.05it/s]\n",
      "epoch 95 iter 16: train loss 0.73854. lr 3.235605e-04: 100%|██████████| 17/17 [00:04<00:00,  4.04it/s]\n",
      "epoch 96 iter 16: train loss 0.73887. lr 3.188598e-04: 100%|██████████| 17/17 [00:04<00:00,  4.02it/s]\n",
      "epoch 97 iter 16: train loss 0.72368. lr 3.141544e-04: 100%|██████████| 17/17 [00:04<00:00,  3.97it/s]\n",
      "epoch 98 iter 16: train loss 0.70582. lr 3.094455e-04: 100%|██████████| 17/17 [00:04<00:00,  4.02it/s]\n",
      "epoch 99 iter 16: train loss 0.70190. lr 3.047342e-04: 100%|██████████| 17/17 [00:04<00:00,  3.94it/s]\n",
      "epoch 100 iter 16: train loss 0.68987. lr 3.000218e-04: 100%|██████████| 17/17 [00:04<00:00,  3.94it/s]\n",
      "epoch 101 iter 16: train loss 0.67758. lr 2.953094e-04: 100%|██████████| 17/17 [00:04<00:00,  3.95it/s]\n",
      "epoch 102 iter 16: train loss 0.66017. lr 2.905981e-04: 100%|██████████| 17/17 [00:04<00:00,  4.00it/s]\n",
      "epoch 103 iter 16: train loss 0.65877. lr 2.858892e-04: 100%|██████████| 17/17 [00:04<00:00,  4.00it/s]\n",
      "epoch 104 iter 16: train loss 0.64497. lr 2.811837e-04: 100%|██████████| 17/17 [00:04<00:00,  3.98it/s]\n",
      "epoch 105 iter 16: train loss 0.64179. lr 2.764829e-04: 100%|██████████| 17/17 [00:04<00:00,  3.97it/s]\n",
      "epoch 106 iter 16: train loss 0.63559. lr 2.717879e-04: 100%|██████████| 17/17 [00:04<00:00,  3.97it/s]\n",
      "epoch 107 iter 16: train loss 0.62855. lr 2.670999e-04: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 108 iter 16: train loss 0.61758. lr 2.624199e-04: 100%|██████████| 17/17 [00:04<00:00,  3.94it/s]\n",
      "epoch 109 iter 16: train loss 0.60657. lr 2.577493e-04: 100%|██████████| 17/17 [00:04<00:00,  4.03it/s]\n",
      "epoch 110 iter 16: train loss 0.59996. lr 2.530890e-04: 100%|██████████| 17/17 [00:04<00:00,  3.97it/s]\n",
      "epoch 111 iter 16: train loss 0.58367. lr 2.484404e-04: 100%|██████████| 17/17 [00:04<00:00,  3.95it/s]\n",
      "epoch 112 iter 16: train loss 0.58197. lr 2.438044e-04: 100%|██████████| 17/17 [00:04<00:00,  3.92it/s]\n",
      "epoch 113 iter 16: train loss 0.57091. lr 2.391824e-04: 100%|██████████| 17/17 [00:04<00:00,  3.91it/s]\n",
      "epoch 114 iter 16: train loss 0.57271. lr 2.345753e-04: 100%|██████████| 17/17 [00:04<00:00,  3.97it/s]\n",
      "epoch 115 iter 16: train loss 0.55243. lr 2.299844e-04: 100%|██████████| 17/17 [00:04<00:00,  3.94it/s]\n",
      "epoch 116 iter 16: train loss 0.54761. lr 2.254108e-04: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 117 iter 16: train loss 0.54044. lr 2.208555e-04: 100%|██████████| 17/17 [00:04<00:00,  3.99it/s]\n",
      "epoch 118 iter 16: train loss 0.53038. lr 2.163198e-04: 100%|██████████| 17/17 [00:04<00:00,  4.02it/s]\n",
      "epoch 119 iter 16: train loss 0.53014. lr 2.118048e-04: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 120 iter 16: train loss 0.52053. lr 2.073115e-04: 100%|██████████| 17/17 [00:04<00:00,  3.95it/s]\n",
      "epoch 121 iter 16: train loss 0.51295. lr 2.028411e-04: 100%|██████████| 17/17 [00:04<00:00,  3.95it/s]\n",
      "epoch 122 iter 16: train loss 0.51009. lr 1.983946e-04: 100%|██████████| 17/17 [00:04<00:00,  3.97it/s]\n",
      "epoch 123 iter 16: train loss 0.51131. lr 1.939732e-04: 100%|██████████| 17/17 [00:04<00:00,  3.99it/s]\n",
      "epoch 124 iter 16: train loss 0.49467. lr 1.895780e-04: 100%|██████████| 17/17 [00:04<00:00,  3.97it/s]\n",
      "epoch 125 iter 16: train loss 0.48849. lr 1.852101e-04: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 126 iter 16: train loss 0.47873. lr 1.808704e-04: 100%|██████████| 17/17 [00:04<00:00,  3.92it/s]\n",
      "epoch 127 iter 16: train loss 0.47302. lr 1.765602e-04: 100%|██████████| 17/17 [00:04<00:00,  3.93it/s]\n",
      "epoch 128 iter 16: train loss 0.47510. lr 1.722804e-04: 100%|██████████| 17/17 [00:04<00:00,  4.00it/s]\n",
      "epoch 129 iter 16: train loss 0.46640. lr 1.680321e-04: 100%|██████████| 17/17 [00:04<00:00,  3.97it/s]\n",
      "epoch 130 iter 16: train loss 0.46535. lr 1.638164e-04: 100%|██████████| 17/17 [00:04<00:00,  4.01it/s]\n",
      "epoch 131 iter 16: train loss 0.45578. lr 1.596343e-04: 100%|██████████| 17/17 [00:04<00:00,  3.91it/s]\n",
      "epoch 132 iter 16: train loss 0.45213. lr 1.554869e-04: 100%|██████████| 17/17 [00:04<00:00,  3.91it/s]\n",
      "epoch 133 iter 16: train loss 0.44695. lr 1.513751e-04: 100%|██████████| 17/17 [00:04<00:00,  3.94it/s]\n",
      "epoch 134 iter 16: train loss 0.44066. lr 1.473000e-04: 100%|██████████| 17/17 [00:04<00:00,  3.99it/s]\n",
      "epoch 135 iter 16: train loss 0.42818. lr 1.432625e-04: 100%|██████████| 17/17 [00:04<00:00,  3.93it/s]\n",
      "epoch 136 iter 16: train loss 0.43211. lr 1.392637e-04: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 137 iter 16: train loss 0.42977. lr 1.353046e-04: 100%|██████████| 17/17 [00:04<00:00,  3.97it/s]\n",
      "epoch 138 iter 16: train loss 0.41763. lr 1.313862e-04: 100%|██████████| 17/17 [00:04<00:00,  3.93it/s]\n",
      "epoch 139 iter 16: train loss 0.41813. lr 1.275093e-04: 100%|██████████| 17/17 [00:04<00:00,  3.95it/s]\n",
      "epoch 140 iter 16: train loss 0.41200. lr 1.236750e-04: 100%|██████████| 17/17 [00:04<00:00,  3.91it/s]\n",
      "epoch 141 iter 16: train loss 0.40947. lr 1.198842e-04: 100%|██████████| 17/17 [00:04<00:00,  3.97it/s]\n",
      "epoch 142 iter 16: train loss 0.41038. lr 1.161379e-04: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 143 iter 16: train loss 0.40603. lr 1.124369e-04: 100%|██████████| 17/17 [00:04<00:00,  3.92it/s]\n",
      "epoch 144 iter 16: train loss 0.40523. lr 1.087822e-04: 100%|██████████| 17/17 [00:04<00:00,  3.98it/s]\n",
      "epoch 145 iter 16: train loss 0.39535. lr 1.051747e-04: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 146 iter 16: train loss 0.39754. lr 1.016153e-04: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 147 iter 16: train loss 0.38898. lr 9.810479e-05: 100%|██████████| 17/17 [00:04<00:00,  3.91it/s]\n",
      "epoch 148 iter 16: train loss 0.38628. lr 9.464413e-05: 100%|██████████| 17/17 [00:04<00:00,  3.94it/s]\n",
      "epoch 149 iter 16: train loss 0.38674. lr 9.123415e-05: 100%|██████████| 17/17 [00:04<00:00,  3.91it/s]\n",
      "epoch 150 iter 16: train loss 0.38966. lr 8.787567e-05: 100%|██████████| 17/17 [00:04<00:00,  3.98it/s]\n",
      "epoch 151 iter 16: train loss 0.38590. lr 8.456954e-05: 100%|██████████| 17/17 [00:04<00:00,  3.95it/s]\n",
      "epoch 152 iter 16: train loss 0.37149. lr 8.131657e-05: 100%|██████████| 17/17 [00:04<00:00,  3.92it/s]\n",
      "epoch 153 iter 16: train loss 0.37366. lr 7.811757e-05: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 154 iter 16: train loss 0.36515. lr 7.497331e-05: 100%|██████████| 17/17 [00:04<00:00,  3.91it/s]\n",
      "epoch 155 iter 16: train loss 0.36510. lr 7.188458e-05: 100%|██████████| 17/17 [00:04<00:00,  3.99it/s]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 156 iter 16: train loss 0.36846. lr 6.885214e-05: 100%|██████████| 17/17 [00:04<00:00,  3.97it/s]\n",
      "epoch 157 iter 16: train loss 0.35783. lr 6.587674e-05: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 158 iter 16: train loss 0.36345. lr 6.295911e-05: 100%|██████████| 17/17 [00:04<00:00,  4.01it/s]\n",
      "epoch 159 iter 16: train loss 0.35740. lr 6.009997e-05: 100%|██████████| 17/17 [00:04<00:00,  3.98it/s]\n",
      "epoch 160 iter 16: train loss 0.36017. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.00it/s]\n",
      "epoch 161 iter 16: train loss 0.35203. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  3.98it/s]\n",
      "epoch 162 iter 16: train loss 0.34658. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  3.98it/s]\n",
      "epoch 163 iter 16: train loss 0.35008. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  3.93it/s]\n",
      "epoch 164 iter 16: train loss 0.34701. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.01it/s]\n",
      "epoch 165 iter 16: train loss 0.34820. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.05it/s]\n",
      "epoch 166 iter 16: train loss 0.34178. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  3.99it/s]\n",
      "epoch 167 iter 16: train loss 0.34653. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.03it/s]\n",
      "epoch 168 iter 16: train loss 0.34171. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  3.92it/s]\n",
      "epoch 169 iter 16: train loss 0.34197. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  3.98it/s]\n",
      "epoch 170 iter 16: train loss 0.34081. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 171 iter 16: train loss 0.33828. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.07it/s]\n",
      "epoch 172 iter 16: train loss 0.33962. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.02it/s]\n",
      "epoch 173 iter 16: train loss 0.33878. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  3.99it/s]\n",
      "epoch 174 iter 16: train loss 0.34056. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  3.98it/s]\n",
      "epoch 175 iter 16: train loss 0.33221. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 176 iter 16: train loss 0.33398. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  3.97it/s]\n",
      "epoch 177 iter 16: train loss 0.32910. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  3.92it/s]\n",
      "epoch 178 iter 16: train loss 0.33327. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  3.99it/s]\n",
      "epoch 179 iter 16: train loss 0.32880. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.00it/s]\n",
      "epoch 180 iter 16: train loss 0.32845. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.01it/s]\n",
      "epoch 181 iter 16: train loss 0.32857. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.08it/s]\n",
      "epoch 182 iter 16: train loss 0.33096. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.05it/s]\n",
      "epoch 183 iter 16: train loss 0.32637. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.06it/s]\n",
      "epoch 184 iter 16: train loss 0.32473. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.04it/s]\n",
      "epoch 185 iter 16: train loss 0.33256. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.07it/s]\n",
      "epoch 186 iter 16: train loss 0.32295. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.04it/s]\n",
      "epoch 187 iter 16: train loss 0.32289. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.05it/s]\n",
      "epoch 188 iter 16: train loss 0.31948. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.02it/s]\n",
      "epoch 189 iter 16: train loss 0.32211. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  3.93it/s]\n",
      "epoch 190 iter 16: train loss 0.32283. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  3.94it/s]\n",
      "epoch 191 iter 16: train loss 0.31314. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  3.89it/s]\n",
      "epoch 192 iter 16: train loss 0.31746. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  3.95it/s]\n",
      "epoch 193 iter 16: train loss 0.31409. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  3.96it/s]\n",
      "epoch 194 iter 16: train loss 0.31736. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.00it/s]\n",
      "epoch 195 iter 16: train loss 0.31542. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.08it/s]\n",
      "epoch 196 iter 16: train loss 0.30956. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.05it/s]\n",
      "epoch 197 iter 16: train loss 0.31757. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.08it/s]\n",
      "epoch 198 iter 16: train loss 0.31694. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.07it/s]\n",
      "epoch 199 iter 16: train loss 0.30833. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.07it/s]\n",
      "epoch 200 iter 16: train loss 0.30588. lr 6.000000e-05: 100%|██████████| 17/17 [00:04<00:00,  4.05it/s]\n"
     ]
    }
   ],
   "source": [
    "from mingpt.trainer import Trainer, TrainerConfig\n",
    "\n",
    "# initialize a trainer instance and kick off training\n",
    "tconf = TrainerConfig(max_epochs=200, batch_size=512, learning_rate=6e-4,\n",
    "                      lr_decay=True, warmup_tokens=512*20, final_tokens=200*len(train_dataset)*block_size,\n",
    "                      num_workers=4)\n",
    "trainer = Trainer(model, train_dataset, None, tconf)\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "O God, O God! which is the business so harm!\n",
      "Well, lords, and save yourselves; and no oath to be angry\n",
      "That in their embraces: and, to brave the life\n",
      "We have forgot and bandy as that time\n",
      "Have told me and he bids me for this excellent,\n",
      "Now I would say he looks on the banks\n",
      "And give more strength than a wild and provide\n",
      "A salt that with some friendly vow,\n",
      "That from the reaches of the gain and stop the sleeves\n",
      "Do scope that which He should hide for his guard\n",
      "As miser made thee first way from his holy exercise.\n",
      "\n",
      "BUCKINGHAM:\n",
      "Go, rating to London, with all these woful chances\n",
      "Misthink the king and not be satisfied!\n",
      "\n",
      "Son:\n",
      "Was ever son so rued a father's death?\n",
      "\n",
      "Father:\n",
      "The warn's idle buy and blows: and then to make a\n",
      "fire, sir, I will keep my capss with stars out\n",
      "And safely point of good content.\n",
      "Signior Lucentio, let us hence; good gods rest ourselves:\n",
      "We shall we show her own heaven and the king\n",
      "In me resolved: I have seen a lady's nose\n",
      "That has been blue, but not her eyebrows.\n",
      "\n",
      "First Lady:\n",
      "Hark ye;\n",
      "The queen your mother rounds apace: we shall\n",
      "Present our services to a fine new prince\n",
      "One of these days; and then you'ld wanton with us,\n",
      "If we would have you.\n",
      "\n",
      "Second Lady:\n",
      "She is spread of late\n",
      "Into a goodly bulk: good time encounter her!\n",
      "\n",
      "HENRY BOLINGBROKE:\n",
      "I swear.\n",
      "\n",
      "THOMAS MOWBRAY:\n",
      "And I, to keep all this.\n",
      "\n",
      "HENRY BOLINGBROKE:\n",
      "Sweet peace conduct his sweet soul to the bosom\n",
      "Of good old Abraham! Lords appellants,\n",
      "Your differences shall all rest under gage\n",
      "Till we assign you to your days of trial.\n",
      "\n",
      "DUKE OF YORK:\n",
      "Sweet York, what wilt thou do?\n",
      "Wilt thou not hide the trespass of thine own?\n",
      "Have we more sons? or are we like to have?\n",
      "Is not my teeming date drunk up with time?\n",
      "And wilt thou pluck my fair son from mine age,\n",
      "And rob me of a happy mother's name?\n",
      "Is he not like thee? is he not thine own?\n",
      "\n",
      "DORCAS:\n",
      "Whither?\n",
      "\n",
      "MOPSA:\n",
      "It becomes thy oath full well,\n",
      "Thou to me thy secrets tell.\n",
      "\n",
      "DORCAS:\n",
      "Me too, let me go thither.\n",
      "\n",
      "MOPSA:\n",
      "Or thou goest to the orange or mill.\n",
      "\n",
      "DORCAS:\n",
      "If to either, tho\n"
     ]
    }
   ],
   "source": [
    "# alright, let's sample some character-level shakespear\n",
    "from mingpt.utils import sample\n",
    "\n",
    "context = \"O God, O God!\"\n",
    "x = torch.tensor([train_dataset.stoi[s] for s in context], dtype=torch.long)[None,...].to(trainer.device)\n",
    "y = sample(model, x, 2000, temperature=0.9, sample=True, top_k=5)[0]\n",
    "completion = ''.join([train_dataset.itos[int(i)] for i in y])\n",
    "print(completion)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# well that was fun"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
